{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NNTool API Demonstration\n",
    "\n",
    "This notebook contains a short demonstration on how to use NNTool programatically from Python."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First load a sample network. This is Blazeface from Greenwaves Technologies Github based NNMenu repository. NNMenu contains many preported networks for GAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/GreenWaves-Technologies/blaze_face/raw/a3645c152d34b34ea437d7d21b67f1f7051f18de/model/face_detection_front.tflite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now import the NNMenu API. Most of the supported API is exposed through the NNGraph class. NNGraph is NNTool's internal representation of a Neural Network graph. Any method or property of NNGraph that is documented is an official API. Logging in nntool can be controlled with the standard python logging APIs. The root logger is called 'nntool'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nntool.api import NNGraph\n",
    "from nntool.api.utils import model_settings, quantization_options\n",
    "import logging\n",
    "nntool_log = logging.getLogger('nntool')\n",
    "nntool_log.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we load the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NNGraph.load_graph('face_detection_front.tflite')\n",
    "\n",
    "# Model show returns a table of information on the Graph\n",
    "# print(model.show())\n",
    "\n",
    "# Model draw can open or save a PDF with a visual representation of the graph\n",
    "# model.draw()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define a little dataloader that returns normalized inputs ready for the executer. This could be extended to provide random samples from the full set, etc. IT would also be better to cache the files locally rather than downloading them each time.\n",
    "\n",
    "If you just want to import local data you can use the import_data function or FileImporter dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nntool.api.utils import import_data, FileImporter\n",
    "\n",
    "# data = import_data('path/to/file.jpg', norm_func=lambda x: x/255)\n",
    "# FileImporter.from_wildcard('/path/to/files/*.jpg')\n",
    "# FileImporter.from_wildcards(('/path/to/input1_files/*.jpg', '/path/to/input2_files/*.jpg'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "class GitHubDataLoader():\n",
    "    def __init__(self, max_idx, normalize=True, return_index=False):\n",
    "        self._idx = 0\n",
    "        self._max_idx = max_idx\n",
    "        self._normalize = normalize\n",
    "        self._return_index = return_index\n",
    "        self._filter = None\n",
    "\n",
    "    def _get_url(self, idx):\n",
    "        pass\n",
    "\n",
    "    def _normalize_func(self, val):\n",
    "        pass\n",
    "\n",
    "    def _get_name(self, idx):\n",
    "        return idx\n",
    "\n",
    "    def set_filter(self, labels):\n",
    "        self._filter = labels\n",
    "\n",
    "    def __iter__(self):\n",
    "        self._idx = 0\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        while True:\n",
    "            if self._idx > self._max_idx:\n",
    "                raise StopIteration()\n",
    "            idx = self._idx\n",
    "            label = self._get_name(idx)\n",
    "            self._idx += 1\n",
    "            if self._filter is None or label in self._filter:\n",
    "                break\n",
    "        # print(f\"get {self._get_url(idx)}\")\n",
    "        with requests.get(self._get_url(idx), stream=True) as r:\n",
    "            r.raise_for_status()\n",
    "            r.raw.decode_content = True\n",
    "            image = Image.open(r.raw)\n",
    "            image = image.resize((128, 128))\n",
    "            image.mode = \"RGB\"\n",
    "        val = np.array(image, dtype=np.int8).transpose((2, 0, 1))\n",
    "        if self._normalize:\n",
    "            val = self._normalize_func(val)\n",
    "        if self._return_index:\n",
    "            return label, [val]\n",
    "        return [val]\n",
    "\n",
    "class BlazeFaceDataLoader(GitHubDataLoader):\n",
    "    def __init__(self, last=130, **kwargs):\n",
    "        super().__init__(last, **kwargs)\n",
    "\n",
    "    def _normalize_func(self, val):\n",
    "        return (val - 128)/128\n",
    "\n",
    "    def _get_name(self, idx):\n",
    "        return f'{idx:04d}.pgm'\n",
    "\n",
    "    def _get_url(self, idx):\n",
    "        return f\"https://github.com/GreenWaves-Technologies/blaze_face/raw/a3645c152d34b34ea437d7d21b67f1f7051f18de/eval_dataset/{self._get_name(idx)}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The equivalent of the adjust command\n",
    "model.adjust_order()\n",
    "\n",
    "# The equivalent of the fusions --scale8 command. The fusions method can be given a series of fusions to apply\n",
    "# fusions('name1', 'name2', etc)\n",
    "model.fusions('scaled_match_group')\n",
    "\n",
    "# draw the model here again to see the adjusted and fused graph\n",
    "# model.draw()\n",
    "\n",
    "# Lets load an image and execute the graph in float on the normalized data\n",
    "data = next(BlazeFaceDataLoader())\n",
    "\n",
    "# The executer returns all the layer output. Each layer output is an array of the outputs from each output of a layer\n",
    "# Generally layers have one output but some (like a split for example) can have multiple outputs\n",
    "# Here we select the first output of the last layer which in a graph with one output will always be the the\n",
    "# graph output\n",
    "model.execute(data)[-1][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's quantize the graph using a small amount of sample data. You may want to use more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "statistics = model.collect_statistics(BlazeFaceDataLoader(last=3))\n",
    "# The resulting statistics contain detailed information for each layer\n",
    "statistics['input_1']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the statistics to quantize the graph and execute it. The quantization information is saved in the model. The quantize option on execute quantizes the imput data. The dequantize option dequantizes the output after execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quantize the model. quantization options can be supplied for a layer or for the whole model\n",
    "model.quantize(statistics, schemes=['scaled'])\n",
    "data = next(BlazeFaceDataLoader())\n",
    "\n",
    "# Now execute the quantized graph outputing quantized values\n",
    "print(\"execute model without dequantizing data\")\n",
    "print(model.execute(data, quantize=True)[-1][0])\n",
    "\n",
    "# Now execute the graph twice with float and quantized versions and compare the results\n",
    "print(\"execute model comparing float and quantized execution and showing Cosine Similarity\")\n",
    "cos_sim = model.cos_sim(model.execute(data), model.execute(data, quantize=True, dequantize=True))\n",
    "print(cos_sim)\n",
    "# the step idx can be used to index the model to find the layer with the worst cos_sim\n",
    "model[np.argmin(cos_sim)]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets look at how we can compress parameters of the graph using the LUT compressor in GAP9.\n",
    "\n",
    "First we must create a Validator that validates the compressed output of the graph. In this BlazeFace case we are going to validate the QSNR of the output of the graph. The validator may be called many times by the AutoCompressor so we will cache the results of the uncompressed execution of the graph.\n",
    "\n",
    "This is a somewhat artificial example since we are using the QSNR of the output to validate the graph. Normally the validator should use the actual labels or if audio, PESQ, etc. to validate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LAST_IDX = 5\n",
    "\n",
    "outputs = {}\n",
    "for filename, data in BlazeFaceDataLoader(LAST_IDX, return_index=True):\n",
    "    outputs[filename] = model.execute(data)[-1][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now the validator class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nntool.utils.validation_utils import ValidateBase, ValidationResultBase\n",
    "from nntool.api.utils import qsnrs\n",
    "\n",
    "class QSNRResult(ValidationResultBase):\n",
    "    def __init__(self, qsnr, label, margin):\n",
    "        self._qsnr = qsnr\n",
    "        self._label = label\n",
    "        self._margin = margin\n",
    "\n",
    "    @property\n",
    "    def validated(self):\n",
    "        return self._margin >= 0\n",
    "\n",
    "    @property\n",
    "    def margin(self):\n",
    "        return self._margin\n",
    "\n",
    "class QSNRValidator(ValidateBase):\n",
    "\n",
    "    def __init__(self, outputs, min_qsnr):\n",
    "        self._outputs = outputs\n",
    "        self._qsnr = min_qsnr\n",
    "\n",
    "    def _validate(self,\n",
    "                  input_tensors,\n",
    "                  output_tensors,\n",
    "                  input_name):\n",
    "        qsnr = qsnrs([self._outputs[input_name]], output_tensors[-1])[0]\n",
    "        if qsnr > 100:\n",
    "            margin = 1.0\n",
    "        else:\n",
    "            margin = (qsnr - self._qsnr)/self._qsnr\n",
    "        return QSNRResult(qsnr, input_name, margin)        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run the compressor. Since it can take a little time there is a progress function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nntool.api.compression import AutoCompress, print_progress\n",
    "\n",
    "TARGET_QSNR = 15\n",
    "\n",
    "autocompress = AutoCompress(\n",
    "    model,\n",
    "    BlazeFaceDataLoader(LAST_IDX, return_index=True),\n",
    "    QSNRValidator(outputs, TARGET_QSNR))\n",
    "autocompress.tune_all(model.all_constants, progress=print_progress)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading and executing a graph\n",
    "\n",
    "Lets look at how we can load a graph and run on GAP. First lets retrieve a network from NNMENU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/GreenWaves-Technologies/image_classification_networks/raw/50ad1beb9ac784b1f5f3574beb2a4c39a46b2fbc/models/tflite_models/mobilenet_v1_1_0_224_quant.tflite\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now load into nntool and process the graph. For this to complete successfully the GAP_SDK_HOME environment variable must point to your GAP SDK directory. This can be set either in your shell startup script or in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = NNGraph.load_graph(\"mobilenet_v1_1_0_224_quant.tflite\", load_quantization=True)\n",
    "G.adjust_order()\n",
    "G.fusions('scaled_match_group')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can exectute the graph on GVSOC. This method creates a project, builds it and executes it. It can parse some of the output including Autotiler output and perfromance data. The GAP_SDK must be sourced to use this API. It is normal that it takes some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = G.execute_on_target(pretty=True, at_log=True, at_loglevel=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the graph run has finished the res object will contain the requested information. Run the cell below to see the performance information as a pretty table (due to the pretty command above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(res.performance)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's execute the same model again but changing some quantization and model generation settings. The nntool API defines some helper functions to create the settings dictionaries. In this case we target the graph on the NE16. The GVSOC simulation of the NE16 is quite slow so we enable output to see execution progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "G.quantize(graph_options=quantization_options(use_ne16=True))\n",
    "G.adjust_order()\n",
    "res = G.execute_on_target(\n",
    "    pretty=True,\n",
    "    at_log=True,\n",
    "    at_loglevel=1,\n",
    "    print_output=True,\n",
    "    settings=model_settings(l1_size=128000, l2_size=1000000, graph_trace_exec=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving a graph\n",
    "\n",
    "The intermediate graph state inside NNTool can be saved and reloaded using NNGraph.write_graph_state and\n",
    "NNGraph.read_graph_state. The write process also returns a string containing a function that can be used to recreate the graph and a tensors dictionary that it requires as a parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "import os\n",
    "notebook_dir = os.getcwd()\n",
    "graph_file = os.path.join(notebook_dir, \"mygraph.zip\")\n",
    "graph_function, tensors = G.write_graph_state(graphpath=graph_file)\n",
    "with open(os.path.join(notebook_dir, \"graph_function.py\"), 'w') as file:\n",
    "    file.write(graph_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the saved graph can be read back with NNGraph.read_graph_state. The graph state file is actually a python zip module containing the function and tensors file so it can also be imported. Finally the returned function was saved in a file in the cell above and we can import the creation function from that file.\n",
    "\n",
    "The creation function is a good example on how a synthetic graph can be created using the NNGraph API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create using NNGraph API\n",
    "new_graph = NNGraph.read_graph_state(graph_file)\n",
    "new_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create using written graph function\n",
    "from graph_function import create_graph\n",
    "new_graph = create_graph(tensors)\n",
    "new_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create by importing\n",
    "import sys\n",
    "sys.path.insert(0, graph_file)\n",
    "from mygraph import mygraph\n",
    "mygraph"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c2859481f2c8de7f07eb9ae202a57839ab4182f99bcf7ae1efb35a84004b024e"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
